import numpy
import re
import datetime
import time
import pandas as pd
import json
import os

import numpy as np

# // Obstacle type
ObstacleType={
    0:'Invalid',
    1:'Vehicle',
    2:'Pedestrian',
    3:'Rider',
    4:'Traffic_cone',
    5:'Animal',
    6:'Road_debris',
    7:'Fence',
    'Vehicle':{
        0:'unknown',
        1:'bus',
        2:'car',
        3:'truck',
        4:'special',
        5:'tiny',
        6:'van'
    },
    'Rider':{
        0: 'unknown',
        1: 'cyclist',
        2: 'motorcyclist',
        3: 'tricyclist'
    },
    'Pedestrian':{
        0:'unknown',
        1:'adult',
        2:'child'
    }
}

def get_type(type_data):
    name=[]
    for type in type_data:
        name.append(ObstacleType[ObstacleType[type[0]]][type[1]])
    return name

def get_json_data(label_file):
    # 读取预测数据
    with open(label_file, 'r') as f:
        data_label = json.load(f)
    name = [i['obj_type'].lower() for i in data_label]
    xyz = np.array(
        [[i['psr']['position']['x'], i['psr']['position']['y'], i['psr']['position']['z']] for i in data_label])
    lwh = np.array([[i['psr']['scale']['x'], i['psr']['scale']['y'], i['psr']['scale']['z']] for i in data_label])
    r = np.array([[i['psr']['rotation']['z']] for i in data_label])
    box = np.concatenate([xyz, lwh, r], axis=-1)
    # box[:, 2] = box[:, 2] - 3.15
    velocity = np.zeros([name.__len__(), 5]).tolist()
    acceleration = np.zeros([name.__len__(), 5]).tolist()
    if 'gt' not in label_file:
        score = [i['score'] for i in data_label]
        return name, box, velocity, acceleration, score
    else:
        return name, box, velocity, acceleration

pr_data_path='/media/king/8EDA0BA1DA0B84A5/tmp/highway/label/'
pr_json_path='/media/king/8EDA0BA1DA0B84A5/tmp/highway/pr_label_json/'
# path='/home/king/workplace/data/data28_2021_04_03/'

file_name_list=os.listdir(pr_data_path)
count_f=0
for file_name in file_name_list:
    # 设置文件路径
    pr_data_path0 = os.path.join(pr_data_path,file_name)
    pr_json_path0 = os.path.join(pr_json_path,file_name)

    name_pr, box_pr, velocity_pr, acceleration_pr, score = get_json_data(pr_data_path0)

    data_temp = {
        'point_cloud': {
            'frame_id': count_f,
            'next_frame': count_f+1,
            'previous_frame': count_f-1,
            'point_file': -1,
            'offset_z': 1},
        'image_file': file_name[:-5]+'.jpg',
        'radar_file': file_name[:-5]+'.bin',
        'muti_sensor_cfg': 0,
        'annos': {
            'name': name_pr,  # * str < ----- Check:-1 | shape: 1
            'gt_boxes_3d': box_pr.tolist(),  # c <= 2 < ----- Check: -1 | shape: 2
            'num_lidar_pts': -1,  # < ----- Check:-1 | shape: 1x
            'track_id': np.linspace(0,name_pr.__len__()-1,name_pr.__len__()).tolist(),  # < ----- Check:-1 | shape: 2
            'velocity': velocity_pr,  # float / -1 < ----- Check: -1 | shape: 2
            'acceleration': acceleration_pr,  # float / -1 < ----- Check: -1 | shape: 2
            'confidence':score,
            'others': 0
        }
    }
    count_f += 1
    with open(pr_json_path0, 'w') as f:
        json.dump(data_temp, f)

# datas=np.concatenate(datas)
# datas=pd.DataFrame(datas)
# datas.columns=head
# N=5